from typing import Dict, List, Optional, Tuple
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import joblib
from datetime import datetime
import os

class MLPredictor:
    """机器学习预测器"""
    
    def __init__(self, model_dir: str = 'models'):
        """初始化机器学习预测器"""
        self.model_dir = model_dir
        self.models = {}
        self.scalers = {}
        self.feature_columns = None
        os.makedirs(model_dir, exist_ok=True)
        
    def _get_available_features(self, df: pd.DataFrame) -> List[str]:
        """获取可用的特征列"""
        # 基础价格特征
        price_features = ['open', 'high', 'low', 'close', 'volume']
        
        # 技术指标特征
        tech_features = [
            'rsi_14',  # RSI
            'macd', 'macd_signal', 'macd_hist',  # MACD
            'upper_band', 'middle_band', 'lower_band',  # 布林带
            'atr',  # ATR
            'stoch_k', 'stoch_d',  # KD
            'adx',  # ADX
            'mom', 'roc'  # 动量
        ]
        
        # 移动平均线特征
        ma_features = [
            'ma_5', 'ma_10', 'ma_20', 'ma_50', 'ma_200',  # 简单移动平均
            'ema_5', 'ema_10', 'ema_20',  # 指数移动平均
            'volume_sma_20', 'volume_ema_20'  # 成交量移动平均
        ]
        
        # 自定义特征
        custom_features = [
            'price_momentum',  # 价格动量
            'volatility',  # 波动率
            'trend_strength',  # 趋势强度
            'price_position',  # 价格位置
            'volume_trend',  # 成交量趋势
            'rsi_trend'  # RSI趋势
        ]
        
        # 获取所有可用的特征列
        all_features = price_features + tech_features + ma_features + custom_features
        
        # 返回实际存在的特征列
        available_features = [col for col in all_features if col in df.columns]
        return available_features
        
    def prepare_features(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.Series]:
        """准备特征数据"""
        # 确保数据不为空
        if df.empty:
            raise ValueError("输入数据为空")
            
        # 获取可用特征
        available_features = self._get_available_features(df)
        if not available_features:
            raise ValueError("没有找到任何可用的特征")
            
        print("可用特征:", available_features)
        print("数据形状:", df.shape)
        
        # 创建目标变量（未来价格变化）
        df['target'] = df['close'].shift(-1) / df['close'] - 1
        
        # 处理无穷大和异常值
        df = df.replace([np.inf, -np.inf], np.nan)
        
        # 使用前向填充处理缺失值
        df = df.ffill()
        
        # 使用后向填充处理剩余的缺失值
        df = df.bfill()
        
        # 如果还有任何NaN值，用0填充
        df = df.fillna(0)
        
        # 对每一列，将超过3个标准差的值限制在这个范围内
        for col in df.columns:
            if df[col].dtype in [np.float64, np.int64]:
                mean = df[col].mean()
                std = df[col].std()
                df[col] = df[col].clip(mean - 3*std, mean + 3*std)
        
        print("清理后的数据形状:", df.shape)
        
        # 确保所有必需的特征都存在，如果不存在则用0填充
        required_features = [
            'open', 'high', 'low', 'close', 'volume',
            'rsi_14', 'macd', 'macd_signal', 'macd_hist',
            'upper_band', 'middle_band', 'lower_band',
            'atr', 'stoch_k', 'stoch_d', 'adx', 'mom', 'roc',
            'ma_5', 'ma_10', 'ma_20', 'ma_50', 'ma_200',
            'ema_5', 'ema_10', 'ema_20',
            'volume_sma_20', 'volume_ema_20',
            'price_momentum', 'volatility', 'trend_strength', 'price_position', 'volume_trend', 'rsi_trend'
        ]
        
        for feature in required_features:
            if feature not in df.columns:
                df[feature] = 0
        
        # 保存特征列名
        self.feature_columns = required_features
        
        X = df[required_features]
        y = df['target']
        
        # 最后检查确保没有NaN值
        if X.isna().any().any() or y.isna().any():
            raise ValueError("数据清理后仍然存在NaN值")
        
        return X, y
        
    def train_models(self, df: pd.DataFrame, test_size: float = 0.2) -> Dict:
        """训练预测模型"""
        X, y = self.prepare_features(df)
        
        # 数据分割
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, shuffle=False
        )
        
        # 数据标准化
        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)
        
        results = {}
        
        # 训练随机森林模型
        rf_model = RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )
        rf_model.fit(X_train_scaled, y_train)
        
        # 保存模型和scaler
        self.models['rf'] = rf_model
        self.scalers['rf'] = scaler
        
        # 评估随机森林模型
        rf_train_pred = rf_model.predict(X_train_scaled)
        rf_test_pred = rf_model.predict(X_test_scaled)
        
        results['rf'] = {
            'train_mse': mean_squared_error(y_train, rf_train_pred),
            'test_mse': mean_squared_error(y_test, rf_test_pred),
            'train_r2': r2_score(y_train, rf_train_pred),
            'test_r2': r2_score(y_test, rf_test_pred)
        }
        
        # 训练梯度提升模型
        gb_model = GradientBoostingRegressor(
            n_estimators=100,
            max_depth=5,
            random_state=42
        )
        gb_model.fit(X_train_scaled, y_train)
        
        # 保存模型和scaler
        self.models['gb'] = gb_model
        self.scalers['gb'] = scaler
        
        # 评估梯度提升模型
        gb_train_pred = gb_model.predict(X_train_scaled)
        gb_test_pred = gb_model.predict(X_test_scaled)
        
        results['gb'] = {
            'train_mse': mean_squared_error(y_train, gb_train_pred),
            'test_mse': mean_squared_error(y_test, gb_test_pred),
            'train_r2': r2_score(y_train, gb_train_pred),
            'test_r2': r2_score(y_test, gb_test_pred)
        }
        
        return results
        
    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        """使用训练好的模型进行预测"""
        if not self.models or not self.scalers or not self.feature_columns:
            raise ValueError("模型未训练")
            
        # 准备特征
        X, _ = self.prepare_features(df)
        
        if not all(col in X.columns for col in self.feature_columns):
            raise ValueError("输入数据缺少必要的特征")
            
        # 使用相同的特征顺序
        X = X[self.feature_columns]
        
        predictions = {}
        for model_name, model in self.models.items():
            # 标准化数据
            X_scaled = self.scalers[model_name].transform(X)
            predictions[f'{model_name}_pred'] = model.predict(X_scaled)
            
        return pd.DataFrame(predictions, index=df.index)
